import tushare as ts
import pandas as pd
import talib
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
import numpy as np
# 初始化Tushare Pro API
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()

# 股票代码和时间范围
ts_code = '600519.SH'  # 贵州茅台
start_date = '20180101'  # 近5年的起始日期
end_date = '20230101'  # 近5年的结束日期

# 获取股票行情数据
df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
df = df.sort_values(by='trade_date')
df = df.reset_index(drop=True)

# 将字符串类型的日期转换为datetime类型
df['trade_date'] = pd.to_datetime(df['trade_date'])

# 计算常见技术指标
# 简单移动平均线（SMA）
df['SMA_5'] = talib.SMA(df['close'].astype(float).values, timeperiod=5)
df['SMA_10'] = talib.SMA(df['close'].astype(float).values, timeperiod=10)

# 相对强弱指标（RSI）
df['RSI_14'] = talib.RSI(df['close'].astype(float).values, timeperiod=14)

# 布林带（Bollinger Bands）
upperband, middleband, lowerband = talib.BBANDS(df['close'].astype(float).values, timeperiod=20)
df['Upper_Band'] = upperband
df['Middle_Band'] = middleband
df['Lower_Band'] = lowerband
# 打印结果
print(df[['trade_date', 'open', 'high', 'low', 'close',
          'SMA_5', 'SMA_10', 'RSI_14', 'Upper_Band', 'Middle_Band', 'Lower_Band']])
# 计算收益率和目标标签
df['return'] = df['close'].pct_change()
df['target'] = np.where(df['return'].shift(-1) > 0, 1, 0)

# 去除包含NaN值的行
df = df.dropna()

# 选择特征和目标变量
features = ['SMA_5', 'SMA_10', 'RSI_14', 'Upper_Band', 'Middle_Band', 'Lower_Band']
X = df[features]
y = df['target']

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 构建随机森林模型
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# 进行预测
y_pred = model.predict(X_test)

# 模型评价
accuracy = accuracy_score(y_test, y_pred)
print(f"模型准确率: {accuracy}")
print("分类报告:")
print(classification_report(y_test, y_pred))

# 计算基于预测的收益率
# 计算基于预测的收益率
test_df = df.iloc[-len(X_test):].copy()
test_df['predicted'] = y_pred
test_df['strategy_return'] = test_df['predicted'] * test_df['return']
# 使用 iloc 按位置索引获取最后一个元素
cumulative_strategy_return = (1 + test_df['strategy_return']).cumprod().iloc[-1] - 1
print(f"策略累计收益率: {cumulative_strategy_return}")
